#!/bin/bash
set -e
RANK_SIZE=$1
DEVICE_START_ID=$2

EXEC_PATH=$(pwd)
export GLOG_v=3
# export ASCEND_SLOG_PRINT_TO_STDOUT=1
test_dist_8pcs()
{
    export RANK_TABLE_FILE=/user/config/nbstart_hccl.json
    export RANK_SIZE=8
}

test_dist_4pcs()
{
    export RANK_TABLE_FILE=/home/ma-user/work/liulei/ice_pred_0407/ice-predict/gc-ice-pred/rank_table_4pcs.json
    export RANK_SIZE=4
}

test_dist_${RANK_SIZE}pcs
root_dir=test_dist_${RANK_SIZE}pcs
for((i=$DEVICE_START_ID;i<$[$RANK_SIZE+$DEVICE_START_ID];i++))
do
	rm -rf ${root_dir}/device$i
	mkdir -p ${root_dir}/device$i
	cp ./main.py ./${root_dir}/device$i
	cp -r ./src ./${root_dir}/device$i
	# cp -r ./mindearth ./device$i
    cp ./GraphCast.yaml ./${root_dir}/device$i
	cd ./${root_dir}/device$i
	mkdir logs
	export DEVICE_ID=$i
	export RANK_ID=$[$i-$DEVICE_START_ID]
	
	echo "start training for device $i"
	env > env$i.log
	nohup python -u main.py --device_id $i\
		--distribute 1\
		--save_graphs 0\
	    >train${i}.log 2>&1 &
	cd ../../
done